from random_forest import *
from load_data import *
import itertools
import sys
import warnings
from sklearn.ensemble import RandomForestRegressor

warnings.filterwarnings("ignore")

def run():
    train_x, train_y, val_x, val_y = load_data()
    n_estimators_all = [1, 2, 10, 20]
    max_depth_all = [-1, 3, 10]
    min_samples_split_all = [3, 6]
    min_samples_leaf_all = [2, 5]
    min_split_gain_all = [0.0, 0.1]
    feature_rate_all = [0.2, 0.5, 0.8]
    item_rate_all = [0.2, 0.5, 0.8]

    max_r = -100
    for (n_estimators, max_depth, min_samples_split, min_samples_leaf, min_split_gain, feature_rate, item_rate) in \
            itertools.product(n_estimators_all, max_depth_all, min_samples_split_all, min_samples_leaf_all, min_split_gain_all, feature_rate_all, item_rate_all):
        print('------parameters: ', n_estimators, max_depth, min_samples_split, min_samples_leaf, min_split_gain, feature_rate, item_rate)
        rf = RandomForest_Classifier(
                n_estimators=n_estimators, max_depth=max_depth, min_samples_split=min_samples_split,
                min_samples_leaf=min_samples_leaf, min_split_gain=min_split_gain, 
                item_rate=item_rate, feature_rate=feature_rate, random_state=2 
        )
        rf.fit(train_x, train_y)
        print('训练集..............')
        train_pred = rf.predict(train_x)
        train_r = r2_score(train_y, train_pred)
        print(train_r)
        print('预测集..............')
        val_pred = rf.predict(val_x)
        val_r = r2_score(val_y, val_pred)
        print(val_r)

        if val_r > max_r:
            max_r = val_r
            save_model(rf, 'save/Random_Forest.pickle')
    clf = RandomForestRegressor(n_estimators=20)
    # print(train_y)
    clf.fit(train_x, train_y) 
    ori, data = load_test_data()
    prediction = clf.predict(data)
    ori['charges'] = prediction
    ori.to_csv('data/submission.csv', index=False)
    # break

def test():
    rf = load_model('save/Random_Forest.pickle')
    ori, data = load_test_data()
    prediction = rf.predict(data.values.tolist())
    ori['charges'] = prediction
    ori.to_csv('data/submission.csv', index=False)


def cal_score(labels, predictions):
    pass

if __name__=='__main__':
    run()
    # test()
